package cas.ihep.alluxio.impl;

import alluxio.AlluxioURI;
import alluxio.client.block.BlockMasterClient;
import alluxio.client.block.policy.BlockLocationPolicy;
import alluxio.client.block.policy.options.GetWorkerOptions;
import alluxio.client.block.stream.BlockInStream;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.options.InStreamOptions;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.exception.AlluxioException;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.PreconditionMessage;
import alluxio.exception.status.NotFoundException;
import alluxio.exception.status.UnavailableException;
import alluxio.master.block.BlockId;
import alluxio.proto.dataserver.Protocol;
import alluxio.resource.CloseableResource;
import alluxio.wire.BlockInfo;
import alluxio.wire.BlockLocation;
import alluxio.wire.WorkerNetAddress;
import cas.ihep.alluxio.AlluxioBlock;
import cas.ihep.alluxio.AlluxioFile;
import cas.ihep.alluxio.AlluxioSystem;
import com.google.common.base.Preconditions;
import com.google.common.io.ByteStreams;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.InvocationTargetException;
import java.util.List;

public class AlluxioFileIn extends AlluxioBaseFile {

    InStreamOptions mInOptions;
    private OutStreamOptions mOutOptions;

    AlluxioFileIn(AlluxioBaseSystem sys, AlluxioURI path) {
        super(sys, path);
        mInOptions=InStreamOptions.defaults();
        mOutOptions=OutStreamOptions.defaults();
    }

    @Override
    public AlluxioFile open() throws IOException {
        try{
            mStatus=mSys.getFileSystem().getStatus(mUri);
            mBlockSize=mStatus.getBlockSizeBytes();
            mFileLength=mStatus.getLength();
            return this;
        }catch (AlluxioException e){
            throw new IOException(e);
        }
    }

    @Override
    public AlluxioBlock getBlock(int idx) throws IOException{
        long id=mStatus.getBlockIds().get(idx);
        if (mBlocks.containsKey(id)){
            return mBlocks.get(id);
        }else{
            AlluxioBaseBlock blk = findBlock(id);
            mCloser.register(blk);
            mBlocks.put(id,blk);
            return blk;
        }
    }

    @Override
    public AlluxioBlock nextBlock() throws IOException {
        throw new IOException("Unsupported operation");
    }

    @Override
    public boolean canRead() {
        return true;
    }

    @Override
    public boolean canWrite() {
        return false;
    }

    private AlluxioBaseBlock findBlock(long blockId)throws IOException{
        try {
            AlluxioInternal handler=mSys.getInternalHandler();
            FileSystemContext context=mSys.getContext();
            Protocol.OpenUfsBlockOptions openUfsBlockOptions=null;
            if (mStatus.isPersisted()) {
                long blockStart = BlockId.getSequenceNumber(blockId) * mBlockSize;
                openUfsBlockOptions =
                        Protocol.OpenUfsBlockOptions.newBuilder().setUfsPath(mStatus.getUfsPath())
                                .setOffsetInFile(blockStart).setBlockSize(getBlockSize(blockStart))
                                .setMaxUfsReadConcurrency(mInOptions.getMaxUfsReadConcurrency())
                                .setNoCache(!mInOptions.getAlluxioStorageType().isStore())
                                .setMountId(mStatus.getMountId()).build();
            }
            BlockInfo blockInfo;
            try (CloseableResource<BlockMasterClient> masterClientResource =
                         context.acquireBlockMasterClientResource()) {
                blockInfo = masterClientResource.get().getBlockInfo(blockId);
            }
            List<BlockLocation> locations=blockInfo.getLocations();
            if (locations.isEmpty() && openUfsBlockOptions == null) {
                throw new NotFoundException("Block " + blockId + " is unavailable in both Alluxio and UFS.");
            }
            long blockSize=blockInfo.getLength();
            for (BlockLocation location : locations) {
                WorkerNetAddress workerNetAddress = location.getWorkerAddress();
                if (workerNetAddress.getHost().equals(AlluxioSystem.localHostName)) {
                    return new AlluxioReadableBlock(this,workerNetAddress,blockId,blockSize);
                }
            }
            BlockInStream.BlockInStreamSource source = BlockInStream.BlockInStreamSource.UFS;
            WorkerNetAddress address;
            if (locations.isEmpty()) {
                BlockLocationPolicy blockLocationPolicy =
                        Preconditions.checkNotNull(mInOptions.getUfsReadLocationPolicy(),
                                PreconditionMessage.UFS_READ_LOCATION_POLICY_UNSPECIFIED);
                address = blockLocationPolicy
                        .getWorker(GetWorkerOptions.defaults().setBlockWorkerInfos(handler.getWorkerInfoList())
                                .setBlockId(blockId).setBlockSize(blockInfo.getLength()));
                if (address == null) {
                    throw new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage());
                }
            } else {
                address = locations.get(handler.nextInt(locations.size())).getWorkerAddress();
                source = BlockInStream.BlockInStreamSource.REMOTE;
            }
            Protocol.ReadRequest.Builder builder = Protocol.ReadRequest.newBuilder().setBlockId(blockId)
                    .setPromote(mInOptions.getAlluxioStorageType().isPromote());
            if (openUfsBlockOptions != null) {
                builder.setOpenUfsBlockOptions(openUfsBlockOptions);
            }
            try (InputStream is = handler.createNettyBlockInStream(context, address, source,
                    builder.buildPartial(), blockSize, mInOptions)){
                WorkerNetAddress localWorker = context.getLocalWorker();
                if (localWorker != null) {
                    try (OutputStream os =
                                 handler.getOutStream(blockId, blockSize, localWorker, mOutOptions)) {
                        ByteStreams.copy(is,os);
                    }
                    return new AlluxioReadableBlock(this,localWorker,blockId,blockSize);
                }
                throw new IOException("Cannot find block on local worker");
            }
        } catch (IllegalAccessException | InvocationTargetException e){
            throw new IOException(e);
        }
    }
}
